import os
import json
from datetime import datetime

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.pipeline import Pipeline
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import joblib


MODELS_DIR = "models"
ASSETS_DIR = "assets"

def ensure_dirs():
    os.makedirs(MODELS_DIR, exist_ok=True)
    os.makedirs(ASSETS_DIR, exist_ok=True)

def load_data():
    cali = fetch_california_housing(as_frame=True)
    df = cali.frame.copy()
    feature_names = cali.feature_names  # ['MedInc', 'HouseAge', ...]
    target_name = cali.target_names[0]  # 'MedHouseVal'
    X = df[feature_names]
    y = df[target_name]
    return X, y, feature_names, target_name

def build_pipeline(random_state=42):
    # 强力树模型，免标准化，适合非线性
    model = HistGradientBoostingRegressor(
        learning_rate=0.1,
        max_depth=None,
        max_iter=300,
        min_samples_leaf=20,
        l2_regularization=0.0,
        random_state=random_state
    )
    pipe = Pipeline([
        ("model", model)
    ])
    return pipe

def plot_pred_vs_true(y_true, y_pred, path):
    plt.figure(figsize=(6, 6))
    plt.scatter(y_true, y_pred, s=8, alpha=0.5, edgecolors="none")
    lims = [min(y_true.min(), y_pred.min()), max(y_true.max(), y_pred.max())]
    plt.plot(lims, lims, "r--", linewidth=1)
    plt.xlabel("True Median House Value")
    plt.ylabel("Predicted Median House Value")
    plt.title("Predicted vs. True")
    plt.tight_layout()
    plt.savefig(path, dpi=150)
    plt.close()

def plot_residuals(y_true, y_pred, path):
    residuals = y_true - y_pred
    plt.figure(figsize=(6, 4))
    plt.hist(residuals, bins=50, color="#4C72B0", alpha=0.85)
    plt.xlabel("Residuals (True - Pred)")
    plt.ylabel("Count")
    plt.title("Residuals Distribution")
    plt.tight_layout()
    plt.savefig(path, dpi=150)
    plt.close()

def plot_feature_importance(model, feature_names, path):
    # HistGradientBoostingRegressor 提供 feature_importances_
    if hasattr(model, "feature_importances_"):
        importances = model.feature_importances_
        idx = np.argsort(importances)
        plt.figure(figsize=(7, 5))
        plt.barh(np.array(feature_names)[idx], importances[idx], color="#55A868")
        plt.xlabel("Importance")
        plt.title("Feature Importances")
        plt.tight_layout()
        plt.savefig(path, dpi=150)
        plt.close()

def main():
    ensure_dirs()
    X, y, feature_names, target_name = load_data()

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    pipe = build_pipeline()
    pipe.fit(X_train, y_train)
    y_pred = pipe.predict(X_test)

    rmse = float(np.sqrt(mean_squared_error(y_test, y_pred)))
    mae = float(mean_absolute_error(y_test, y_pred))
    r2 = float(r2_score(y_test, y_pred))

    metrics = {
        "timestamp": datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
        "n_train": int(X_train.shape[0]),
        "n_test": int(X_test.shape[0]),
        "features": feature_names,
        "target": target_name,
        "metrics": {
            "RMSE": rmse,
            "MAE": mae,
            "R2": r2
        },
        "model": {
            "type": "HistGradientBoostingRegressor",
            "sklearn_version": __import__("sklearn").__version__
        }
    }

    # 保存模型与元数据
    model_path = os.path.join(MODELS_DIR, "california_housing_pipeline.joblib")
    joblib.dump(pipe, model_path)
    with open(os.path.join(MODELS_DIR, "feature_names.json"), "w", encoding="utf-8") as f:
        json.dump(feature_names, f, ensure_ascii=False, indent=2)
    with open(os.path.join(MODELS_DIR, "metrics.json"), "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)

    # 示例输入
    sample = X_test.iloc[0].to_dict()
    with open(os.path.join(MODELS_DIR, "sample_input.json"), "w", encoding="utf-8") as f:
        json.dump(sample, f, ensure_ascii=False, indent=2)

    # 可视化
    plot_pred_vs_true(
        y_test.to_numpy(),
        y_pred,
        os.path.join(ASSETS_DIR, "pred_vs_true.png")
    )
    plot_residuals(
        y_test.to_numpy(),
        y_pred,
        os.path.join(ASSETS_DIR, "residuals.png")
    )
    # 取出底层模型
    base_model = pipe.named_steps["model"]
    plot_feature_importance(
        base_model, feature_names, os.path.join(ASSETS_DIR, "feature_importance.png")
    )

    print("Training complete.")
    print(f"RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}")
    print(f"Saved model to: {model_path}")
    print("Plots saved in assets/")

if __name__ == "__main__":
    main()